<#
    .DESCRIPTION
        Runbook to initiate and monitor malware scans on Azure Storage Accounts with a specific tag using Interactive Azure login.
#>

[CmdletBinding()]
param(
    [Parameter(Mandatory = $true)]
    [string]$SubscriptionId,

    [Parameter(Mandatory = $true)]
    [string]$TagKey = "ScanForMalware",

    [Parameter(Mandatory = $true)]
    [string]$TagValue = "True",

    [Parameter(Mandatory = $true)]
    [string]$OutputCsvPath
)

function Write-ErrorAndExit {
    param([string]$ErrorMessage)
    Write-Error $ErrorMessage
    throw $ErrorMessage
}

function Connect-ToAzure {
    param(
        [string]$SubscriptionId
    )
    try {
        Write-Output "Logging in to Azure using Interactive login..."
        # This will prompt for login via a browser
        Connect-AzAccount -ErrorAction Stop

        # Set the Azure context to the specified Subscription
        Set-AzContext -SubscriptionId $SubscriptionId -ErrorAction Stop
        Write-Output "Azure context set to Subscription ID: $SubscriptionId"
    }
    catch {
        Write-ErrorAndExit "Failed to set Azure context. Error: $_"
    }
}

function Start-MalwareScan {
    param(
        [Microsoft.Azure.Commands.Management.Storage.Models.PSStorageAccount]$StorageAccount,
        [string]$SubscriptionId
    )
    $resourceGroupName = $StorageAccount.ResourceGroupName
    $storageAccountName = $StorageAccount.StorageAccountName
    $uri = "https://management.azure.com/subscriptions/$SubscriptionId/resourceGroups/$resourceGroupName/providers/Microsoft.Storage/storageAccounts/$storageAccountName/providers/Microsoft.Security/defenderForStorageSettings/current/startMalwareScan?api-version=2024-10-01-preview"

    try {
        $response = Invoke-AzRestMethod -Method POST -Uri $uri -ErrorAction Stop
        return $response
    }
    catch {
        Write-Warning "Failed to start scan for '$storageAccountName'. Error: $_"
        return $null
    }
}

function Get-MalwareScanStatus {
    param(
        [Microsoft.Azure.Commands.Management.Storage.Models.PSStorageAccount]$StorageAccount,
        [string]$SubscriptionId
    )
    $resourceGroupName = $StorageAccount.ResourceGroupName
    $storageAccountName = $StorageAccount.StorageAccountName
    $uri = "https://management.azure.com/subscriptions/$SubscriptionId/resourceGroups/$resourceGroupName/providers/Microsoft.Storage/storageAccounts/$storageAccountName/providers/Microsoft.Security/defenderForStorageSettings/current/malwareScans/latest?api-version=2024-10-01-preview"

    try {
        $response = Invoke-AzRestMethod -Method GET -Uri $uri -ErrorAction Stop
        return $response
    }
    catch {
        Write-Warning "Failed to get scan status for '$storageAccountName'. Error: $_"
        return $null
    }
}

try {
    # Authenticate and set the subscription context
    Connect-ToAzure -SubscriptionId $SubscriptionId

    Write-Output "`nFetching storage accounts with tag '$TagKey=$TagValue'..."
    $storageAccounts = Get-AzStorageAccount | Where-Object { $_.Tags[$TagKey] -eq $TagValue }
    if (-not $storageAccounts) {
        Write-ErrorAndExit "No storage accounts found with tag '$TagKey=$TagValue'."
    }

    Write-Output "Found $($storageAccounts.Count) storage account(s) with the specified tag.`n"
    $scanStatuses = @{}
    $scanResults = @()

    foreach ($storageAccount in $storageAccounts) {
        Write-Output "Starting malware scan for storage account: $($storageAccount.StorageAccountName)"
        $response = Start-MalwareScan -StorageAccount $storageAccount -SubscriptionId $SubscriptionId
        if ($response) {
            $responseContent = $response.Content | ConvertFrom-Json
            if ($responseContent.scanStatus -in @("Queued", "WaitingForCompletion")) {
                Write-Output "Successfully initiated scan for '$($storageAccount.StorageAccountName)'. Status: $($responseContent.scanStatus)`n"
                $scanStatuses[$storageAccount.StorageAccountName] = @{
                    Status      = $responseContent.scanStatus
                    LastChecked = Get-Date
                    ScanDetails = $responseContent
                }
            }
            else {
                Write-Warning "Failed to start scan for '$($storageAccount.StorageAccountName)'. Status: $($responseContent.scanStatus)`n"
                $scanStatuses[$storageAccount.StorageAccountName] = @{
                    Status      = "Failed"
                    LastChecked = Get-Date
                    ScanDetails = $responseContent
                }
            }
        }
        else {
            Write-Warning "Failed to start scan for '$($storageAccount.StorageAccountName)'.`n"
            $scanStatuses[$storageAccount.StorageAccountName] = @{
                Status      = "Failed"
                LastChecked = Get-Date
                ScanDetails = $null
            }
        }
    }

    # Continuously check the status of each scan until all are completed or failed
    $allCompleted = $false
    while (-not $allCompleted) {
        $allCompleted = $true
        foreach ($accountName in $scanStatuses.Keys) {
            $currentStatus = $scanStatuses[$accountName]
            if ($currentStatus.Status -notin @("Completed", "Failed")) {
                $allCompleted = $false
                $storageAccount = $storageAccounts | Where-Object { $_.StorageAccountName -eq $accountName }
                $status = Get-MalwareScanStatus -StorageAccount $storageAccount -SubscriptionId $SubscriptionId
                if ($status) {
                    $statusContent = $status.Content | ConvertFrom-Json
                    $scanStatuses[$accountName].Status = $statusContent.scanStatus
                    $scanStatuses[$accountName].LastChecked = Get-Date
                    $scanStatuses[$accountName].ScanDetails = $statusContent
                    $blobSummary = $statusContent.scanSummary.blobs
                    Write-Output "Storage Account: $accountName"
                    Write-Output " Status: $($statusContent.scanStatus)"
                    Write-Output " Total Blobs Scanned: $($blobSummary.totalBlobsScanned)"
                    Write-Output " Malicious Blobs Count: $($blobSummary.maliciousBlobsCount)"
                    Write-Output " Scanned Blobs in GB: $([math]::Round($blobSummary.scannedBlobsInGB, 4))`n"
                }
                else {
                    Write-Warning "Failed to get scan status for '$accountName'.`n"
                }
            }
        }

        if (-not $allCompleted) {
            Write-Output "Waiting 10 seconds before next status check...`n"
            Start-Sleep -Seconds 10
        }
    }

    # Collect scan results for CSV export
    $overallSummary = @{
        TotalStorageAccounts   = $storageAccounts.Count
        SuccessfulScans        = 0
        FailedScans            = 0
        TotalBlobsScanned      = 0
        TotalMaliciousBlobs    = 0
        TotalSkippedBlobs      = 0
        TotalScannedBlobsInGB  = 0.0
        EstimatedTotalScanCost = 0.0
    }

    foreach ($accountName in $scanStatuses.Keys) {
        $status = $scanStatuses[$accountName]
        Write-Output "----------------------------------------"
        Write-Output "Storage Account: $accountName"
        Write-Output "Status: $($status.Status)"
        Write-Output "Last Checked: $($status.LastChecked)"
        
        if ($status.Status -eq "Completed" -and $status.ScanDetails) {
            $details = $status.ScanDetails
            Write-Output " Scan ID: $($details.scanId)"
            Write-Output " Scan Start Time: $($details.scanStartTime)"
            Write-Output " Scan End Time: $($details.scanEndTime)"
            Write-Output " Total Blobs Scanned: $($details.scanSummary.blobs.totalBlobsScanned)"
            Write-Output " Malicious Blobs Count: $($details.scanSummary.blobs.maliciousBlobsCount)"
            Write-Output " Skipped Blobs Count: $($details.scanSummary.blobs.skippedBlobsCount)"
            Write-Output " Scanned Blobs in GB: $($details.scanSummary.blobs.scannedBlobsInGB)"
            Write-Output " Estimated Scan Cost (USD): $($details.scanSummary.estimatedScanCostUSD)`n"
            
            # Update overall summary
            $overallSummary.SuccessfulScans++
            $overallSummary.TotalBlobsScanned += $details.scanSummary.blobs.totalBlobsScanned
            $overallSummary.TotalMaliciousBlobs += $details.scanSummary.blobs.maliciousBlobsCount
            $overallSummary.TotalSkippedBlobs += $details.scanSummary.blobs.skippedBlobsCount
            $overallSummary.TotalScannedBlobsInGB += $details.scanSummary.blobs.scannedBlobsInGB
            $overallSummary.EstimatedTotalScanCost += $details.scanSummary.estimatedScanCostUSD

            # Add to the CSV results
            $scanResults += [PSCustomObject]@{
                StorageAccountName    = $accountName
                Status                = $status.Status
                LastChecked           = $status.LastChecked
                ScanId                = $details.scanId
                ScanStartTime         = $details.scanStartTime
                ScanEndTime           = $details.scanEndTime
                TotalBlobsScanned     = $details.scanSummary.blobs.totalBlobsScanned
                MaliciousBlobsCount   = $details.scanSummary.blobs.maliciousBlobsCount
                SkippedBlobsCount     = $details.scanSummary.blobs.skippedBlobsCount
                ScannedBlobsInGB      = $details.scanSummary.blobs.scannedBlobsInGB
                EstimatedScanCostUSD  = $details.scanSummary.estimatedScanCostUSD
            }
        }
        elseif ($status.Status -eq "Failed") {
            Write-Output " Scan not started or failed.`n"
            $overallSummary.FailedScans++
        }
        else {
            Write-Output " Scan Status: $($status.Status)`n"
        }
    }

    # Export results to CSV
    Write-Output "Exporting scan results to CSV..."
    $scanResults | Export-Csv -Path $OutputCsvPath -NoTypeInformation
    Write-Output "Results exported to $OutputCsvPath"

    # Display the overall summary
    Write-Output "----------------------------------------"
    Write-Output "`nOverall Summary:"
    Write-Output " Total Storage Accounts: $($overallSummary.TotalStorageAccounts)"
    Write-Output " Successful Scans: $($overallSummary.SuccessfulScans)"
    Write-Output " Failed Scans: $($overallSummary.FailedScans)"
    Write-Output " Total Blobs Scanned: $($overallSummary.TotalBlobsScanned)"
    Write-Output " Total Malicious Blobs: $($overallSummary.TotalMaliciousBlobs)"
    Write-Output " Total Skipped Blobs: $($overallSummary.TotalSkippedBlobs)"
    Write-Output " Total Scanned Blobs in GB: $([math]::Round($overallSummary.TotalScannedBlobsInGB, 4))"
    Write-Output " Estimated Total Scan Cost (USD): $([math]::Round($overallSummary.EstimatedTotalScanCost, 6))"
}
catch {
    Write-ErrorAndExit "An error occurred: $_"
}
finally {
    Write-Output "`nScript execution completed."
}
